﻿#pragma once

#include "frontend/token.h"

namespace mspc {

// 前向声明
struct ASTNode;
struct File;
struct ImportStmt;
struct StructDef;
struct MessageDef;
struct FieldDef;
struct EnumDef;
struct EnumValue;
struct OptionDef;
class ASTVisitor;

// AST节点基类
struct ASTNode {
    virtual ~ASTNode() = default;
    size_t line = 1;
    size_t column = 1;

    // 访问者模式支持
    virtual void accept(class ASTVisitor& visitor) = 0;
};

// 导入语句
struct ImportStmt : ASTNode {
    std::string filename;
    bool is_public = false;

    void accept(ASTVisitor& visitor) override;
};

// 字段选项
struct OptionDef : ASTNode {
    std::string name;
    std::string value;

    OptionDef() = default;
    OptionDef(const std::string& n, const std::string& v)
        : name(n)
        , value(v) {}

    void accept(ASTVisitor& visitor) override;
};

// 枚举值
struct EnumValue : ASTNode {
    std::string name;
    int32_t value;
    std::vector<std::unique_ptr<OptionDef>> options;

    EnumValue() = default;
    EnumValue(const std::string& n, int32_t v)
        : name(n)
        , value(v) {}

    void accept(ASTVisitor& visitor) override;
};

// 枚举定义
struct EnumDef : ASTNode {
    std::string name;
    std::vector<std::unique_ptr<EnumValue>> values;
    std::vector<std::unique_ptr<OptionDef>> options;

    EnumDef() = default;
    EnumDef(const std::string& n)
        : name(n) {}

    const EnumValue* find_value(const std::string& value_name) const {
        for (const auto& val : values) {
            if (val->name == value_name) return val.get();
        }
        return nullptr;
    }

    const EnumValue* find_value(int32_t value_num) const {
        for (const auto& val : values) {
            if (val->value == value_num) return val.get();
        }
        return nullptr;
    }

    void accept(ASTVisitor& visitor) override;
};

// 字段定义
struct FieldDef : ASTNode {
    std::string rule;
    std::string type;
    std::string name;
    uint32_t number;
    std::string default_value;
    // std::vector<std::unique_ptr<OptionDef>> options;

    FieldDef() = default;
    FieldDef(const std::string& r, const std::string& t, const std::string& n, uint32_t num)
        : rule(r)
        , type(t)
        , name(n)
        , number(num) {}

    bool is_repeated() const { return rule == "repeated"; }

    /*
    const OptionDef* get_option(const std::string& option_name) const {
        for (const auto& opt : options) {
            if (opt->name == option_name) return opt.get();
        }
        return nullptr;
    }
    */
    void accept(ASTVisitor& visitor) override;
};

// 结构定义
struct StructDef : ASTNode {
    std::string name;
    std::vector<std::unique_ptr<FieldDef>> fields;
    std::vector<std::unique_ptr<OptionDef>> options;

    StructDef() = default;
    StructDef(const std::string& n)
        : name(n) {}

    const FieldDef* find_field(const std::string& field_name) const {
        for (const auto& field : fields) {
            if (field->name == field_name) return field.get();
        }
        return nullptr;
    }

    void accept(ASTVisitor& visitor) override;
};

// 消息定义
struct MessageDef : ASTNode {
    std::string name;
    uint32_t id = 0;
    std::vector<std::unique_ptr<FieldDef>> fields;
    std::vector<std::unique_ptr<OptionDef>> options;
    std::vector<std::unique_ptr<StructDef>> nested_structs;
    std::vector<std::unique_ptr<MessageDef>> nested_messages;
    std::vector<std::unique_ptr<EnumDef>> nested_enums;

    MessageDef() = default;
    MessageDef(const std::string& n)
        : name(n) {}

    const FieldDef* find_field(const std::string& field_name) const {
        for (const auto& field : fields) {
            if (field->name == field_name) return field.get();
        }
        return nullptr;
    }

    const FieldDef* find_field(uint32_t field_number) const {
        for (const auto& field : fields) {
            if (field->number == field_number) return field.get();
        }
        return nullptr;
    }

    const OptionDef* get_option(const std::string& option_name) const {
        for (const auto& opt : options) {
            if (opt->name == option_name) return opt.get();
        }
        return nullptr;
    }

    void accept(ASTVisitor& visitor) override;
};

// 文件节点（根节点）
struct File : ASTNode {
    std::string filename;
    std::string syntax_version = "msp/v1";
    std::string package_name;
    std::vector<std::unique_ptr<ImportStmt>> imports;
    std::vector<std::unique_ptr<StructDef>> structs;
    std::vector<std::unique_ptr<MessageDef>> messages;
    std::vector<std::unique_ptr<EnumDef>> enums;

    File() = default;
    File(const std::string& syntax)
        : syntax_version(syntax) {}

    const StructDef* find_struct(const std::string& struct_name) const {
        for (const auto& str : structs) {
            if (str->name == struct_name) return str.get();
        }
        return nullptr;
    }

    const MessageDef* find_message(const std::string& message_name) const {
        for (const auto& msg : messages) {
            if (msg->name == message_name) return msg.get();
        }
        return nullptr;
    }

    const EnumDef* find_enum(const std::string& enum_name) const {
        for (const auto& enum_def : enums) {
            if (enum_def->name == enum_name) return enum_def.get();
        }
        return nullptr;
    }

    void accept(ASTVisitor& visitor) override;
};

// AST访问者接口
class ASTVisitor {
public:
    virtual ~ASTVisitor() = default;

    virtual void visit(File& file) = 0;
    virtual void visit(ImportStmt& import) = 0;
    virtual void visit(StructDef& struct_def) = 0;
    virtual void visit(MessageDef& message) = 0;
    virtual void visit(FieldDef& field) = 0;
    virtual void visit(EnumDef& enum_def) = 0;
    virtual void visit(EnumValue& enum_value) = 0;
    virtual void visit(OptionDef& option) = 0;
};

// AST节点accept方法实现
inline void ImportStmt::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void OptionDef::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void EnumValue::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void EnumDef::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void FieldDef::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void StructDef::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void MessageDef::accept(ASTVisitor& visitor) { visitor.visit(*this); }
inline void File::accept(ASTVisitor& visitor) { visitor.visit(*this); }

} // namespace mspc